from executors import PyExecutor, ACMExecutor
from generators import PyGenerator, model_factory
from generators.py_generate import get_code_body
from typing import List, Dict, Tuple, Any, Optional, Union
from multiprocessing import Pool
from filelock import FileLock
import random
from transformers import AutoTokenizer
from utils import *
import sys
import json
import os
tokenizer = AutoTokenizer.from_pretrained(os.getenv('TOKENIZER_NAME'))

class Debugger:
    def __init__(self, model_name: str, port: str = "", is_acm: bool = False):
        self.exe = PyExecutor() if not is_acm else ACMExecutor()
        self.gen = PyGenerator()
        self.model = model_factory(model_name, port)
        self.is_acm = is_acm
    
    def prepare_debug_context(self, item: Dict[str, Any], dataset_type: str) -> Tuple[str, List[Any]]:
        """准备调试上下文，包括当前函数实现和测试用例"""
        # 获取当前函数实现
        cur_func_impl = item.get("seed", "")
        
        # 获取测试用例
        tests = item.get("test", [])
        
        # 如果没有测试用例，返回空列表
        if not tests:
            return cur_func_impl, []
        
        # 选择第一个测试用例
        selected_test = tests[0]
        
        # 格式化测试数据
        if dataset_type == "ACM":
            # ACM 模式下的测试数据格式化
            # 检查测试数据的格式并正确获取输入和输出
            if isinstance(selected_test, dict):
                input_data = selected_test.get('input', '')
                output_data = selected_test.get('output', '')
                actual_output = selected_test.get('actual_output', None)
                error = selected_test.get('error', '')
            else:
                # 如果测试数据不是字典格式，尝试解析它
                try:
                    test_data = json.loads(selected_test)
                    input_data = test_data.get('input', '')
                    output_data = test_data.get('output', '')
                    actual_output = test_data.get('actual_output', None)
                    error = test_data.get('error', '')
                except:
                    # 如果无法解析，使用原始测试数据
                    input_data = selected_test
                    output_data = None
                    actual_output = None
                    error = None
            
            # 确保输入和输出都是字符串
            input_str = str(input_data)
            output_str = str(output_data)
            actual_output_str = str(actual_output) if actual_output else 'null'
            error_str = str(error) if error else ''
            
            test_str = f"Input:\n{input_str}\nExpected Output:\n{output_str}"
            if actual_output_str != 'null':
                test_str += f"\nActual Output:\n{actual_output_str}"
            if error_str:
                test_str += f"\nError:\n{error_str}"
        else:
            # 非 ACM 模式下的测试数据格式化
            test_str = f"Input: {selected_test['input']}\nExpected Output: {selected_test['output']}"
            if selected_test.get('actual_output'):
                test_str += f"\nActual Output: {selected_test['actual_output']}"
            if selected_test.get('error'):
                test_str += f"\nError: {selected_test['error']}"
        
        # 返回原始测试用例列表，而不是格式化的测试字符串
        return cur_func_impl, tests
    
    def prepare_debug_function(self, cur_func_impl: str, item: Dict[str, Any], dataset_type: str) -> str:
        """准备用于调试的函数实现，添加必要的注释"""
        if dataset_type == "ACM":
            # ACM 模式下不需要添加注释，直接返回原始代码
            return cur_func_impl
        elif dataset_type in ["HumanEval", "MBPP"]:
            # 添加注释
            if not find_comment(cur_func_impl, item["entry_point"]):
                debug_cur_func_impl = insert_comment(cur_func_impl, extrace_comment(item["prompt"]), item["entry_point"])
            else:
                debug_cur_func_impl = cur_func_impl
        elif dataset_type in ["TransCoder"]:
            # 添加C++翻译作为注释
            debug_cur_func_impl = convert_comment(item["prompt"]) + cur_func_impl
        else:
            debug_cur_func_impl = cur_func_impl
        
        return debug_cur_func_impl
    
    def select_failed_test(self, failed_tests: List[Any]) -> Optional[Any]:
        """随机选择一个失败的测试用例"""
        if not failed_tests:
            return None
            
        if self.is_acm:
            # ACM赛制测试数据
            return failed_tests[random.randint(0, len(failed_tests)-1)]
        else:
            # 原有的断言测试
            return failed_tests[random.randint(0, len(failed_tests)-1)] if len(failed_tests) >= 1 else None
    
    def initialize_messages(self) -> Union[List, str]:
        """根据模型类型初始化消息"""
        return [] if self.model.is_chat else ""
    
    def update_token_count(self, messages: Union[List, str], token_nums: int) -> int:
        """更新token计数"""
        if isinstance(messages, str):
            token_nums += len(tokenizer.tokenize(messages))
        else:
            token_nums += sum([len(tokenizer.tokenize(msg.content)) for msg in messages])
        return token_nums
    
    def debug_iteration(self, item: Dict[str, Any], cur_func_impl: str, tests_i: List[Any], 
                       dataset_type: str, level: str, messages: List[Dict[str, Any]], token_nums: int) -> Tuple[str, List[Dict[str, Any]], int, bool]:
        """执行一次调试迭代"""
        # 准备调试函数
        debug_cur_func_impl = self.prepare_debug_function(cur_func_impl, item, dataset_type)
        
        # 选择一个失败的测试用例
        is_passing, failed_tests, _ = self.exe.execute(cur_func_impl, tests_i)
        if is_passing:
            return cur_func_impl, messages, token_nums, True
        
        selected_test = failed_tests[0]
        
        # 格式化测试数据
        if dataset_type == "ACM":
            # ACM 模式下的测试数据格式化
            # 检查测试数据的格式并正确获取输入和输出
            if isinstance(selected_test, dict):
                input_data = selected_test.get('input', '')
                output_data = selected_test.get('output', '')
                actual_output = selected_test.get('actual_output', None)
                error = selected_test.get('error', '')
            else:
                # 如果测试数据不是字典格式，尝试解析它
                try:
                    test_data = json.loads(selected_test)
                    input_data = test_data.get('input', '')
                    output_data = test_data.get('output', '')
                    actual_output = test_data.get('actual_output', None)
                    error = test_data.get('error', '')
                except:
                    # 如果无法解析，使用原始测试数据
                    input_data = selected_test
                    output_data = None
                    actual_output = None
                    error = None
            
            # 确保输入和输出都是字符串
            input_str = str(input_data)
            output_str = str(output_data)
            actual_output_str = str(actual_output) if actual_output else 'null'
            error_str = str(error) if error else ''
            
            test_str = f"Input:\n{input_str}\nExpected Output:\n{output_str}"
            if actual_output_str != 'null':
                test_str += f"\nActual Output:\n{actual_output_str}"
            if error_str:
                test_str += f"\nError:\n{error_str}"
        else:
            # 非 ACM 模式下的测试数据格式化
            test_str = f"Input: {selected_test['input']}\nExpected Output: {selected_test['output']}"
            if selected_test.get('actual_output'):
                test_str += f"\nActual Output: {selected_test['actual_output']}"
            if selected_test.get('error'):
                test_str += f"\nError: {selected_test['error']}"
        
        # 生成新的实现
        if dataset_type == "ACM":
            # ACM 模式下不需要指定入口函数
            messages = self.gen.ldb_debug(
                item["prompt"], 
                debug_cur_func_impl, 
                test_str, 
                None, 
                self.model, 
                messages, 
                dataset_type, 
                level
            )
            
            # 生成新的函数实现
            messages = self.gen.ldb_generate(
                item["prompt"], 
                self.model, 
                messages, 
                dataset_type
            )
            
            # 从消息中提取函数实现
            if isinstance(messages, list):
                # 如果是消息列表，从最后一条消息中提取函数实现
                for msg in reversed(messages):
                    if msg.role == "assistant":
                        cur_func_impl = get_code_body(msg.content)
                        break
            else:
                # 如果是字符串，直接使用
                cur_func_impl = get_code_body(messages)
        else:
            # 非 ACM 模式下需要指定入口函数
            messages = self.gen.ldb_debug(
                item["prompt"], 
                debug_cur_func_impl, 
                test_str, 
                item["entry_point"], 
                self.model, 
                messages, 
                dataset_type, 
                level
            )
            
            # 生成新的函数实现
            messages = self.gen.ldb_generate(
                item["prompt"], 
                self.model, 
                messages, 
                dataset_type
            )
            
            # 从消息中提取函数实现
            if isinstance(messages, list):
                # 如果是消息列表，从最后一条消息中提取函数实现
                for msg in reversed(messages):
                    if msg.role == "assistant":
                        cur_func_impl = get_code_body(msg.content)
                        break
            else:
                # 如果是字符串，直接使用
                cur_func_impl = get_code_body(messages)
        
        # 更新消息历史和token计数
        if hasattr(self.gen, 'token_nums'):
            token_nums += self.gen.token_nums
        else:
            # 如果PyGenerator没有token_nums属性，则使用默认值0
            token_nums += 0
        
        # 检查新实现
        is_passing, _, _ = self.exe.execute(cur_func_impl, tests_i)
        
        return cur_func_impl, messages, token_nums, is_passing
    
    def debug_task(self, i: int, item: Dict[str, Any], log_path: str, num_items: int, 
                   pass_at_k: int, max_iters: int, level: str = "block") -> None:
        """调试单个任务"""
        cur_pass = 0
        is_solved = False
        implementations = []
        test_feedback = []
        cur_func_impl = ""
        dataset_type = item["task_id"].split("/")[0]
        token_nums = 0
        
        while cur_pass < pass_at_k and not is_solved:
            cur_iter = 0
            messages = []  # 初始化为空列表而不是字符串
            
            # 准备调试上下文
            cur_func_impl, tests_i = self.prepare_debug_context(item, dataset_type)
            implementations.append(cur_func_impl)
            
            # 执行初始实现并获取失败的测试
            is_passing, failed_tests, _ = self.exe.execute(cur_func_impl, tests_i)
            test_feedback.append(failed_tests)
            
            # 如果初始实现通过测试，检查是否解决了问题
            if is_passing:
                is_solved = self.exe.evaluate(item["entry_point"], cur_func_impl, item["test"], timeout=10)
                if is_solved:
                    item["solution"] = cur_func_impl
                break
            
            # 迭代调试直到达到最大迭代次数或解决问题
            while cur_iter < max_iters:
                cur_func_impl, messages, token_nums, is_passing = self.debug_iteration(
                    item, cur_func_impl, tests_i, dataset_type, level, messages, token_nums
                )
                
                implementations.append(cur_func_impl)
                test_feedback.append(failed_tests)
                
                # 如果通过测试或达到最大迭代次数，检查是否解决了问题
                if is_passing or cur_iter == max_iters - 1:
                    if is_passing:
                        print(f'{item["task_id"]} pass generated tests, check real tests')
                    else:
                        print(f'{item["task_id"]} fail generated tests, check real tests')
                    
                    if not is_solved:
                        is_solved = self.exe.evaluate(item["entry_point"], cur_func_impl, item["test"], timeout=10)
                        if is_solved:
                            item["solution"] = cur_func_impl
                    
                    cur_iter += 1
                    sys.stdout.flush()
                    break
                
                cur_iter += 1
                sys.stdout.flush()
            
            cur_pass += 1
        
        # 更新结果
        item["is_passing"] = is_passing
        item["is_solved"] = is_solved
        item["implementations"] = implementations
        item["test_feedback"] = test_feedback
        item["solution"] = cur_func_impl
        item["generated_test"] = tests_i
        item["debug_iter"] = cur_iter
        item["token_nums"] = token_nums
        
        # 保存结果
        with FileLock(log_path + ".lock"):
            write_jsonl(log_path, [item], append=True)
        
        print(f'completed {i+1}/{num_items}')

    def debug_single_problem(self, problem_data: Dict[str, Any], max_iters: int = 5) -> Dict[str, Any]:
        """调试单个问题"""
        # 初始化结果
        result = {
            "task_id": problem_data["task_id"],
            "prompt": problem_data["prompt"],
            "solution": "",
            "debug_history": [],
            "test_results": [],
            "token_usage": 0,
            "is_acm": problem_data.get("is_acm", False),
            "is_solved": False  # 添加 is_solved 字段
        }
        
        # 准备调试上下文
        dataset_type = problem_data["task_id"].split("/")[0]
        cur_func_impl, tests_i = self.prepare_debug_context(problem_data, dataset_type)
        
        # 初始化调试状态
        cur_iter = 0
        messages = []  # 初始化为空列表而不是字符串
        token_nums = 0
        
        # 开始调试迭代
        while cur_iter < max_iters:
            # 执行一次调试迭代
            cur_func_impl, messages, token_nums, is_passing = self.debug_iteration(
                problem_data, cur_func_impl, tests_i, dataset_type, "block", messages, token_nums
            )
            
            # 更新结果
            result["debug_history"].append({
                "iteration": cur_iter,
                "code": cur_func_impl,
                "is_passing": is_passing
            })
            
            # 如果所有测试都通过，则结束调试
            if is_passing:
                result["solution"] = cur_func_impl
                result["token_usage"] = token_nums
                result["is_solved"] = True  # 设置 is_solved 为 True
                break
            
            # 增加迭代次数
            cur_iter += 1
        
        # 如果达到最大迭代次数仍未解决，则使用最后一次的代码作为解决方案
        if not result["solution"]:
            result["solution"] = cur_func_impl
            result["token_usage"] = token_nums
        
        return result

def debug(i, item, log_path, model_name, num_items, pass_at_k, max_iters, port="", level="block", is_acm=False):
    """调试单个任务的入口函数"""
    debugger = Debugger(model_name, port, is_acm)
    debugger.debug_task(i, item, log_path, num_items, pass_at_k, max_iters, level)

def run_ldb(
    dataset: List[dict],
    model_name: str,
    max_iters: int,
    n_proc: int,
    pass_at_k: int,
    log_path: str,
    verbose: bool,
    seedfile: str = None,
    testfile: str = None,
    port: str = "",
    level: str = "block",
    is_acm: bool = False
) -> None:
    print("Number of proc:", n_proc)
    num_items = len(dataset)
    args = iter([(i, item, log_path, model_name, num_items, pass_at_k, max_iters, port, level, is_acm) for i, item in enumerate_resume(dataset, log_path, seedfile, testfile)])
    if n_proc == 1:
        for item in args:
            debug(*item)
    else:
        with Pool(n_proc) as pool:
            pool.starmap(debug, args)
    print("Accuracy:", count_solved(log_path))
    
class DebugVisualizer:
    """调试过程可视化类"""
    
    def __init__(self, debug_result: Dict[str, Any]):
        self.debug_result = debug_result
        self.is_acm = debug_result.get("is_acm", False)
    
    def get_iteration_summary(self, iteration: int) -> Dict[str, Any]:
        """获取指定迭代的摘要信息"""
        if iteration < 0 or iteration >= len(self.debug_result["debug_history"]):
            return {"error": "Invalid iteration index"}
            
        step = self.debug_result["debug_history"][iteration]
        return {
            "iteration": step["iteration"],
            "code_snippet": self._get_code_snippet(step["implementation"]),
            "failed_tests": self._format_failed_tests(step["failed_tests"]),
            "is_passing": step["is_passing"]
        }
    
    def get_final_result(self) -> Dict[str, Any]:
        """获取最终结果摘要"""
        return {
            "is_solved": self.debug_result.get("is_solved", False),
            "final_solution": self._get_code_snippet(self.debug_result.get("solution", "")),
            "token_usage": self.debug_result.get("token_usage", 0),
            "total_iterations": len(self.debug_result.get("debug_history", []))
        }
    
    def get_debug_timeline(self) -> List[Dict[str, Any]]:
        """获取调试时间线"""
        timeline = []
        for i, step in enumerate(self.debug_result["debug_history"]):
            timeline.append({
                "iteration": i,
                "status": "solved" if step["is_passing"] else "failed",
                "code_snippet": self._get_code_snippet(step["code"]),
                "failed_test_count": len(step["failed_tests"]) if "failed_tests" in step and step["failed_tests"] else 0
            })
        return timeline
    
    def _get_code_snippet(self, code: str, max_lines: int = 10) -> str:
        """获取代码片段，限制行数"""
        lines = code.split("\n")
        if len(lines) <= max_lines:
            return code
        return "\n".join(lines[:max_lines]) + "\n..."
    
    def _format_failed_tests(self, failed_tests: List[Any]) -> List[Dict[str, Any]]:
        """格式化失败的测试用例"""
        if not failed_tests:
            return []
            
        formatted_tests = []
        for test in failed_tests:
            if self.is_acm:
                # 检查测试数据的格式
                if isinstance(test, dict):
                    # 如果已经是字典格式，直接使用
                    formatted_tests.append({
                        "input": test.get('input', {}),
                        "expected": test.get('output', {}),
                        "actual": test.get('actual_output', None),
                        "error": test.get('error', '')
                    })
                else:
                    # 如果不是字典格式，尝试解析它
                    try:
                        test_data = json.loads(test)
                        formatted_tests.append({
                            "input": test_data.get('input', {}),
                            "expected": test_data.get('output', {}),
                            "actual": test_data.get('actual_output', None),
                            "error": test_data.get('error', '')
                        })
                    except:
                        # 如果无法解析，创建一个简单的测试数据
                        formatted_tests.append({
                            "input": test,
                            "expected": None,
                            "actual": None,
                            "error": "Failed to parse test data"
                        })
            else:
                formatted_tests.append({
                    "test_case": test
                })
        return formatted_tests
    
def debug_acm_example():
    """
    展示如何使用修改后的代码进行ACM模式的调试
    
    示例数据结构:
    {
        "task_id": "leetcode/1",
        "prompt": "def twoSum(nums: List[int], target: int) -> List[int]:",
        "entry_point": "twoSum",
        "test": [
            {
                "input": {"nums": [2, 7, 11, 15], "target": 9},
                "output": [0, 1]
            },
            {
                "input": {"nums": [3, 2, 4], "target": 6},
                "output": [1, 2]
            }
        ],
        "seed": "def twoSum(nums: List[int], target: int) -> List[int]:\n    # 错误的实现\n    for i in range(len(nums)):\n        for j in range(i+1, len(nums)):\n            if nums[i] + nums[j] == target:\n                return [i, j]\n    return []",
        "is_acm": True
    }
    """
    # 创建调试器实例
    debugger = Debugger(model_name="gpt-4", is_acm=True)
    
    # 准备问题数据
    problem_data = {
        "task_id": "leetcode/1",
        "prompt": "def twoSum(nums: List[int], target: int) -> List[int]:",
        "entry_point": "twoSum",
        "test": [
            {
                "input": {"nums": [2, 7, 11, 15], "target": 9},
                "output": [0, 1]
            },
            {
                "input": {"nums": [3, 2, 4], "target": 6},
                "output": [1, 2]
            }
        ],
        "seed": """def twoSum(nums: List[int], target: int) -> List[int]:
    # 错误的实现
    for i in range(len(nums)):
        for j in range(i+1, len(nums)):
            if nums[i] + nums[j] == target:
                return [i, j]
    return []""",
        "is_acm": True
    }
    
    # 执行调试
    debug_result = debugger.debug_single_problem(problem_data, max_iters=5)
    
    # 创建可视化器
    visualizer = DebugVisualizer(debug_result)
    
    # 获取调试时间线
    timeline = visualizer.get_debug_timeline()
    for step in timeline:
        print(f"Iteration {step['iteration']}: {step['status']}")
        print(f"Failed tests: {step['failed_test_count']}")
        print(f"Code snippet:\n{step['code_snippet']}\n")
    
    # 获取最终结果
    final_result = visualizer.get_final_result()
    print(f"Problem solved: {final_result['is_solved']}")
    print(f"Token usage: {final_result['token_usage']}")
    print(f"Final solution:\n{final_result['final_solution']}")
    
    return debug_result
    